package com.bigdata.flink.watermark;

import com.bigdata.flink.util.StreamEnvUtil;
import org.apache.flink.api.common.eventtime.SerializableTimestampAssigner;
import org.apache.flink.api.common.eventtime.Watermark;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.KeyedStream;
import org.apache.flink.streaming.api.datastream.SideOutputDataStream;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.windowing.ProcessWindowFunction;
import org.apache.flink.streaming.api.windowing.assigners.TumblingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.util.Collector;
import org.apache.flink.util.OutputTag;

import java.time.Duration;
import java.util.Date;

// 事件时间、水位线处理
// 具有事件时间的特殊数据类型，用来衡量窗口数据的处理进度
public class WatermarkEventTimeTest {
  public static void main(String[] args) throws Exception {

    StreamExecutionEnvironment env = StreamEnvUtil.getEnv();
    env.setParallelism(3);

    DataStreamSource<String> socketSource = StreamEnvUtil.getSocketSource(env);
    SingleOutputStreamOperator<String> operator = socketSource.assignTimestampsAndWatermarks(getWatermarkStrategy());

    // 多并行度算子，三个map算子，执行后合并到keyBy的同一算子
    // 合并时，水位线已最小的map算子水位线进行传递
    // 如果取最小的，但是某个算子一直没数据 那就没有水位线传递，则设置超时等待水位线时间
    SingleOutputStreamOperator<String> map = operator.map(f -> String.valueOf(f));

    KeyedStream<String, String> keyedStream = map.keyBy(f -> f.split("\\s")[0]);

    OutputTag<String> deadData = new OutputTag<>("deadData", Types.STRING);
    SingleOutputStreamOperator<String> streamOperator = keyedStream
        .window(TumblingEventTimeWindows.of(Time.seconds(10)))
        // forBoundedOutOfOrderness(Duration.ofSeconds(3)) 水位线迟到数据，可以延迟关窗，关窗时触发计算
        // allowedLateness 迟到数据再次延迟关窗，不过数据来一条触发一条计算
        .allowedLateness(Time.seconds(2))
        // 如果迟到数据之后，还有该窗口的乱序数据到来，则输出到侧输出流
        .sideOutputLateData(deadData)
        .process(new ProcessWindowFunction<String, String, String, TimeWindow>() {

          @Override
          public void process(String string, ProcessWindowFunction<String, String, String, TimeWindow>.Context context, Iterable<String> elements, Collector<String> out) throws Exception {
            long start = context.window().getStart();
            Date date1 = new Date(start);
            long end = context.window().getEnd();
            Date date2 = new Date(end);
            out.collect(String.format("key=%s, 窗口大小[%s  %s], 元素个数%s, %s", string, date1, date2, elements.spliterator().estimateSize(), elements));
          }
        });

    streamOperator.print();
    SideOutputDataStream<String> sideOutput = streamOperator.getSideOutput(deadData);
    sideOutput.print("deadData");

    env.execute();
  }


  // 水位线提取，水位线 = 数据最大时间 - 延迟时间 -1ms, 当水位线 >= 窗口关闭时间则触发窗口关闭
  // string数据格式 s1 1
  public static WatermarkStrategy<String> getWatermarkStrategy() {
    return WatermarkStrategy.<String>forBoundedOutOfOrderness(Duration.ofSeconds(3))
        // 多并行度空闲等待时长
        .withIdleness(Duration.ofSeconds(10))
        .withTimestampAssigner((SerializableTimestampAssigner<String>) (element, recordTimestamp) -> Integer.parseInt(element.split("\\s")[1]) * 1000L);
  }
}
